/******************************************************************************

                                COST.C

  Title:        Source code for cost functions w/o constraints.
  Author:       Hanna Szoke
  Date:         6/10/00 (constraint costs separated off)
  Function:     This module contains functions used by Solve and Back
		for calculating the various hkl cost functions.

*******************************************************************************/

#include "speden.h"	/* includes util.h, cellparams.h, symmetry.h, dims.h */
#include "speden_Np.h"		/* Np space arrays also in back.c, qshare.c */
#include "cost.h"		/* for cost function calculation */
#include "sp.h"

COMPLEX *Oar ;          /* array for holding FFT transform of np's */
real	*Phases ;	/* for holding Oar phases in prep_grad	*/
static	int	first = TRUE ;
real	*ReList ;
real	*phext_expfac ;
real	ph_coef ;

static	real	e_funct_hkl(real *, real *) ;
static	real	prep_grad() ;
static	real	prep_phgrad() ;
static	real	e_robust_funct_hkl(real *, real *) ;
static	void	compute_grad(real *) ;
void	convolve(int, real *, real *, COMPLEX *) ;
void	deconvolve(int, real *, real *, COMPLEX *) ;



void e_funct(real *csol, real *ressq, real *grad, int *iqflag)

/*
real	*csol ;		current solution array	
real	*ressq ;	(returned) residue  
real	*grad ;		(returned) gradient array
int	*iqflag ;	(returned) flag re success	
*/

{
     real	e_funct_np(real *, real, int) ;
     real	e_funct_cs(real *, real *, real) ;

     real	*pcsol, *ptotp, *psnump, *pknownp ;
     real	con_cost[MAXCONSTR] ;
     real	*pg ;
     real	lamcoef ;
     int	i, n, target_number ;

     if (first) {
	  first = FALSE ;
	  
          if (robust) {
	       sprintf(message, 
	       "Using the robust hkl cost function with coefficient %g.", 
	       rob_factor) ;
	       printTwice(message) ;
          }
          else
	       printTwice("Using the original hkl cost function.") ;
     }

     /************************************************************
     Initialize constraint costs.
     Initialize gradient; each partial cost function will 
     accumulate into the gradient, so that ordering is arbitrary.
     ************************************************************/

     for (i = 0; i < MAXCONSTR; i++)
          con_cost[i] = 0 ;

     for (n = 0, pg = grad; n < Npextended; n++, pg++)
          *pg = 0 ;

     /************************************************************
     Do the basic cost calculation in Fourier space. 
     Then compute physical space cost functions based on the 
     various constraints:
        target(s), crystallographic symmetry, 
        and phase extension (in Fourier space).
     In most cases, the space cost function is applied to the total
     electrons/voxel in the starting model (knownp) + those found 
     in previous iterations (snump) + those found in this iteration 
     (csol).  The exception in e_funct_cs(), for which only csol 
     need be considered.
     ************************************************************/

     if (robust)
	  hkl_cost = e_robust_funct_hkl(csol, grad) ;
     else
	  hkl_cost = e_funct_hkl(csol, grad) ;

     if (Nconstraints > 0) {

	for (n=0, ptotp=totp, pcsol=csol, psnump=snump, pknownp=knownp ;
             n < Nptotal; n++, ptotp++, pcsol++, psnump++, pknownp++)

	     *ptotp = *pcsol + *psnump + *pknownp ;

        for (i = 0, target_number = 0; i < Nconstraints; i++) {

	     lamcoef = relwt_con[i] * con_coef[i] ;

	     switch (con_type[i]) {

	     case TARGET :
		  con_cost[i] = e_funct_np(grad, lamcoef, target_number) ;
		  target_number++ ;		/* .. to select next target */
		  break ;

	     case PHASE_EXT :
                  /*** No longer a real-space target; handled using Nfft ***/  
		  break ;

	     case CS :
		  con_cost[i] = e_funct_cs(csol, grad, lamcoef) ;
		  target_number++ ;		/* .. to select next target */
		  break ;

	     default :
		  break ;

             }	
         }
     }
     /************************************************************
     Summarize the full cost function and report.  Separate terms 
     are reported in the appropriate e_funct_*(). 
     ************************************************************/

        *ressq = hkl_cost ;

        for (i = 0; i < Nconstraints; i++)
	     *ressq += con_cost[i] ;

        if (costfile_option) {
             if (Nconstraints > 0) 
	          fprintf(fp_cost, ", Total: %10.4e\n", *ressq) ;
	     else
	          fprintf(fp_cost, "\n") ;
        }

	fflush(fp_cost) ;

	*iqflag = (hkl_cost<= sp_min) ? -1:0 ; 	

	return ;
}

        /***********************************************************************
	This formerly large function has been broken up into three parts:
	convolve() that creates Oar from the input csol; prep_grad() that
	prepares gradient arrays; and compute_grad() that does the deconvolutions.
	real	*csol ;		current solution array
	real	*grad ;		(returned) gradient array
        ***********************************************************************/

real	e_funct_hkl(real *csol, real *grad)
{
	void	hrDoSFT(real *, COMPLEX *) ;

	real	ef_hkl_cost ;
        void	hrGradSFT(COMPLEX *, real *) ;
        real	phext_cost ;

        convolve(Nfft, csol, expfac, Oar) ;

        if (HighRes)
             hrDoSFT(csol+Nptotal, Oar) ;

        ef_hkl_cost =  prep_grad() ;
        phext_cost = (thereis_phext) ? prep_phgrad() : 0 ;
        
        compute_grad(grad) ;
        if (HighRes) 
             hrGradSFT(Oar, grad+Nptotal) ;

	return(ef_hkl_cost + phext_cost) ;
}

         /***********************************************************************
         Replace O by R+O in Oar; calculate phases and hold them in Phases;
	 Make Fclist.  In a loop over list points, calculate resid + (Fc - Fo).
	 Reuse Oar for kernel in deconvolution, weighted resid with Phases.   
        ***********************************************************************/

real	prep_grad()
{
     int	n ;
     real	resid ;
     real	wsq ;
     real	subres = 0 ;
     real	*pF, *wt ;
     real	*pFc, *pP ;
     real	amp0, amp, newamp ;
     COMPLEX	*pOar, *pR ;

     amp0 = Oar->re + R->re ;

     for (n = 0, pOar = Oar, pR = R; n < Nfob; n++, pOar++, pR++) {
         pOar->re += pR->re ;
         pOar->im += pR->im ;

	 getAmpPhase(pOar, amp0, &amp, Phases+n) ;
     }

     for (n = 0; n < Nfob; n++) 
	  *(ReList + n) = Cabs(Oar+n) ;

         /**********************************************************************
         At this point, we can re-use Oar for the gradients. 
         **********************************************************************/

     for (n = 0, pOar = Oar; n < Nfob; n++, pOar++) {
         pOar->re = 0. ;
         pOar->im = 0. ;
     } 

     for (n = 0, pFc=ReList, pF=Folist, wt=hkl_weight; n < Nfob; 
	   n++, pFc++, pF++, wt++) {

         resid = *pFc - *pF ;
         wsq = *wt * *wt ;
         subres += (wsq / 2.) * resid * resid ;

	 newamp = (wsq / 2.) * resid ;

         (Oar + n)->re = newamp ;
         (Oar + n)->im = newamp ; 
     } 

     for (n = 0, pOar = Oar, pP = Phases; n < Nfob; n++, pOar++, pP++) {
         pOar->re *= -cos(*pP*DTOR) ;
         pOar->im *= -sin(*pP*DTOR) ;
     } 

     hkl_cost = subres ;

     if (costfile_option) 
          fprintf(fp_cost, "Hkl space: %10.4e", subres) ;
        
     return (hkl_cost) ;
}

real	prep_phgrad()
{
	int	n, j ;
	real	ressq = 0. ;
	COMPLEX	cdiff ;
	real	phext_cost ;

        /* This code does the same as the comparable code in Back. 
        Note that cdiff is the targetted value (Fc) minus the current
        value from this iteration (Oar) minus those values previously
        found in earlier iterations + those known from the start (R).*/
     
	for (n = 0, j = Nfob; n < Nfc; n++, j++) {
 
	   cdiff.re = *(Fc_relist+n) - (Oar + j)->re - (R + j)->re ;
	   cdiff.im = *(Fc_imlist+n) - (Oar + j)->im - (R + j)->im ;

	   ressq += cdiff.re * cdiff.re + cdiff.im * cdiff.im ; 

	   (Oar + j)->re = cdiff.re ;
	   (Oar + j)->im = cdiff.im ;
        } 

	phext_cost = .5 * ph_coef * ressq ;

	if (costfile_option) 
	     fprintf(fp_cost, ", phext = %10.4e", phext_cost) ;

        return(phext_cost) ;
}

void	compute_grad(real *grad)
{
	int	j ;
	real	*pgrad ;

        /***********************************************************************
        Compute the gradient: g= Re [DFT(-) (Diff * exp(-eta...)) ].
        SUM over the kernel elements that share a resolution before 
	deconvolving.  See Eq. (40), Paper V.
	Note that deconvolution now implies accumulation into the grad array;
	this means that the grad array must be initialized 
	and deconvolve() must accumulate!
        ***********************************************************************/

	for (j = 0, pgrad = grad; j < Nptotal; j++, pgrad++)
	     *pgrad = 0 ;

        deconvolve(Nfft, grad, expfac, Oar) ;

        for (j = 0, pgrad = grad; j < Nptotal; j++, pgrad++)
             *pgrad *= 2. ;
}

real	e_robust_funct_hkl(real *csol, real *grad)  /* currently unused */
{
	int	n ;
	real	dsq, resid ;
	real	Aprime, cprime ;
	real	co_factor ;
	real	der_factor ;
/*	real	step_size_correction ;  Didn't seem to help matters... */
	COMPLEX	exp_phi ;

	COMPLEX	*pOar, *pR ;	
	real	*pFc ;
	real	*pF, *wt ;
	real	subres ;

int N1 = 0, N2 = 0, N3 = 0 ;

        convolve(Nfob, csol, expfac, Oar) ;

        /*********************************************************************
         Replace O by R+O in Oar; calculate |R+O| and hold it in dsq;
	 Put |R+O| - Fobs in resid and collect the weighted residuals.

	 Compute residues (for native and derivatives separately), using a
	 robust cost function that is linear beyond the inner well -

	 f = A' (x/c')^2           if |x| < c'   and
	 f = A' ( 2|x|/c' - 1.)    if |x| > c'
	 
	 where x = resid, c' = Ro*sqrt(norm_fac) / *wt and A' = Ro^2*norm_fac.
	 and Ro is the (input) rob_factor.

         Also, prepare to compute the gradient; exp_phi is (R + O) / |R + O|.
	 We use a constant gradient beyond the inner well.
	 (We must check the denominator and take action if |R+O| is too small.)
	 The F(000) is always handled using the quadratic form!

	 Reuse Oar for kernel in deconvolution, wt^2*(|R+O|-|Fobs|)*exp_phi/2.
	 See Eq. (34) - (39). Paper V.
        *********************************************************************/

        for (n = 0, pOar=Oar, pR=R;  n < Nfob; n++, pOar++, pR++) {
	     pOar->re += pR->re ;
	     pOar->im += pR->im ;
        }

        for (n = 0; n < Nfob; n++) 
	     *(ReList + n) = Cabs(Oar+n) ;

	subres = 0 ;
        for (n = 0, pFc=ReList, pF=Folist, wt=hkl_weight; n < Nfob; 
	     n++, pFc++, pF++, wt++) {

                  dsq = fabs(*pFc) ;
                  resid = dsq - *pF ;
		  cprime = rob_factor * sqrt(norm_fac) / *wt ;
		  Aprime = rob_factor * rob_factor * norm_fac ;
	 
	          if ((resid > cprime && n > 0)) {
		       subres += Aprime * (2 * resid / cprime - 1.) ;
		       co_factor = Aprime /cprime ;
		       der_factor = 1. ;
	/*	       step_size_correction = 2 - cprime / resid ; */
N1++ ;
                  }
	          else if ((resid < -cprime && n > 0)) {
		       subres += Aprime * (-2 * resid / cprime  - 1.) ;
		       co_factor = Aprime /cprime ;
		       der_factor = -1. ;
	/*	       step_size_correction = 2 + cprime / resid ; */
N2++ ;
                  }
	          else {
		       subres += Aprime * (resid/cprime) * (resid/cprime) ;
		       co_factor = Aprime / (cprime * cprime) ;
		       der_factor = resid ;
	/*	       step_size_correction = 1. ; */
N3++ ;
                  }
		  co_factor *= 1. / 2 ;

	          if (*pF > 0) {
		       if (dsq > EPS_PHI) {
		            exp_phi.re = pOar->re / dsq ;
		            exp_phi.im = pOar->im / dsq ;
                       }
		       else {
		            exp_phi.re = (pOar->re >= 0) ? 1. : -1. ;
		            exp_phi.im = 0. ;
                       }
/*		       pOar->re = step_size_correction * 
				  co_factor * der_factor * exp_phi.re ;
		       pOar->im = step_size_correction * 
				  co_factor * der_factor * exp_phi.im ; 
*/
		       pOar->re = co_factor * der_factor * exp_phi.re ;
		       pOar->im = co_factor * der_factor * exp_phi.im ;
                  } 
	          else {
	               pOar->re *= co_factor ;
	               pOar->im *= co_factor ;
                  } 
        } 

	subres *= 1. / 2 ;

        compute_grad(grad) ;

        /* Reports */

	if (costfile_option) 
	    fprintf(fp_cost, "Hkl nat: %10.4e cats: %6i, %6i, %6i", 
		    subres,N1,N2,N3) ;

	return(subres) ;
}

void	Lb_funct(real *nump, real *fval, real *grad, int *iqflag)
{
	int	n ;
	real	ressq = 0. ;
	COMPLEX	cdiff ;

	convolve(Nfc, nump, expfac, Oar) ;

	for (n = 0; n < Nfc; n++) {

	  /* Sign seemed to be wrong so i just swapped terms */
	  /* Swapped again */
	   cdiff.re =   *(Fc_relist+n) - (Oar + n)->re;
	   cdiff.im =   *(Fc_imlist+n) - (Oar + n)->im;

	   ressq += cdiff.re * cdiff.re + cdiff.im * cdiff.im ; 

	   (Oar + n)->re = cdiff.re ;
	   (Oar + n)->im = cdiff.im ;
        } 

	*fval = .5 * ressq ;

	if (costfile_option) {
	     fprintf(fp_cost, "Residual = %g\n", *fval) ;
	     fflush(fp_cost) ;
        }

	*iqflag = 0 ;
	if(ressq <= DISCRP*discrp_frac)        /* DISCRP is defined in speden.h,
						discrp_frac is input var. */
	     *iqflag = -1 ;

/*** Compute the gradient: g(back) = Re [DFT(-) (Diff * exp(-eta...)) ] ***/

        for (n = 0; n < Nptotal; n++)
	     *(grad+n) = 0 ;

        deconvolve(Nfc, grad, expfac, Oar) ;
	/*
        * SKL June 27, 2006
        * Multiplied the gradient by 2
        */
        for (n = 0; n < Nptotal; n++)
         *(grad+n) *= 2. ;

}


	/************************************************
	minimal_Lb_funct() is the list-equivalent of
	(old) minimal_b_funct() for off-grid fcalc input.
	************************************************/

void	minimal_Lb_funct(real *fval)
{
	int	n ;
	real	ressq = 0 ;
        COMPLEX	cdiff ;

	convolve(Nfc, nump, expfac, Oar) ;

	ressq = 0 ;

	for (n = 0; n < Nfc; n++) {

	   cdiff.re =  *(Fc_relist+n) - (Oar + n)->re ;
	   cdiff.im =  *(Fc_imlist+n) - (Oar + n)->im ;

	   ressq += cdiff.re * cdiff.re + cdiff.im * cdiff.im ; 
        } 

	*fval = .5 * ressq ;

	if (costfile_option) {
	     fprintf(fp_cost, "Residual = %g\n", *fval) ;
	     fflush(fp_cost) ;
        }
	return ;
}
